#include <asm.h>
#include <csr.h>

.macro SAVE_CONTEXT
  .local _restore_kernel_tpsp
  .local _save_context
  /*
   * If coming from userspace, preserve the user thread pointer and load
   * the kernel thread pointer.  If we came from the kernel, sscratch
   * will contain 0, and we should continue on the current TP.
   */
  csrrw tp, CSR_SSCRATCH, tp
  bnez tp, _save_context

_restore_kernel_tpsp:
  csrr tp, CSR_SSCRATCH
  sd sp, PCB_KERNEL_SP(tp)
_save_context:
  sd sp, PCB_USER_SP(tp)
  ld sp, PCB_KERNEL_SP(tp)

  /* TODO: save all general purpose registers here! */
  sd zero,OFFSET_REG_ZERO(tp)
  //return address
  sd ra,OFFSET_REG_RA(tp)
  //pointers
  sd sp,OFFSET_REG_SP(tp)
  sd gp,OFFSET_REG_GP(tp)
  sd tp,OFFSET_REG_TP(tp)
  //temporary
  sd t1,OFFSET_REG_T1(tp)
  sd t0,OFFSET_REG_T0(tp)
  sd t2,OFFSET_REG_T2(tp)
  //saved registers
  sd s0,OFFSET_REG_S0(tp)
  sd s1,OFFSET_REG_S1(tp)
  //args
  sd a0,OFFSET_REG_A0(tp)
  sd a1,OFFSET_REG_A1(tp)
  sd a2,OFFSET_REG_A2(tp)
  sd a3,OFFSET_REG_A3(tp)
  sd a4,OFFSET_REG_A4(tp)
  sd a5,OFFSET_REG_A5(tp)
  sd a6,OFFSET_REG_A6(tp)
  sd a7,OFFSET_REG_A7(tp)
  //saved registers
  sd s2,OFFSET_REG_S2(tp)
  sd s3,OFFSET_REG_S3(tp)
  sd s4,OFFSET_REG_S4(tp)
  sd s5,OFFSET_REG_S5(tp)
  sd s6,OFFSET_REG_S6(tp)
  sd s7,OFFSET_REG_S7(tp)
  sd s8,OFFSET_REG_S8(tp)
  sd s9,OFFSET_REG_S9(tp)
  sd s10,OFFSET_REG_S10(tp)
  sd s11,OFFSET_REG_S11(tp)
  //temporary 
  sd t3,OFFSET_REG_T3(tp)
  sd t4,OFFSET_REG_T4(tp)
  sd t5,OFFSET_REG_T5(tp)
  sd t6,OFFSET_REG_T6(tp)
  
  /*
   * Disable user-mode memory access as it should only be set in the
   * actual user copy routines.
   *
   * Disable the FPU to detect illegal usage of floating point in kernel
   * space.
   */
  li t0, SR_SUM | SR_FS

  /* TODO: save sstatus, sepc, stval, scause and sscratch on user stack */
  //privileged
  csrr t3,sstatus
  sd t3,OFFSET_REG_SSTATUS(tp)
  csrr t3,sepc
  sd t3,OFFSET_REG_SEPC(tp)
  
  csrr t3,stval
  sd t3,OFFSET_REG_SBADADDR(tp)
  csrr t3,scause
  sd t3,OFFSET_REG_SCAUSE(tp)
  ld t3,OFFSET_REG_T3(tp)
.endm

.macro RESTORE_CONTEXT
  /* TODO: Restore all registers and sepc,sstatus */
  ld zero,OFFSET_REG_ZERO(tp)
  //return address
  ld ra,OFFSET_REG_RA(tp)
  //pointers
  ld sp,OFFSET_REG_SP(tp)
  ld gp,OFFSET_REG_GP(tp)
  //ld tp,OFFSET_REG_TP(tp)
  //temporary
  ld t1,OFFSET_REG_T1(tp)
  ld t0,OFFSET_REG_T0(tp)
  ld t2,OFFSET_REG_T2(tp)
  //saved registers
  ld s0,OFFSET_REG_S0(tp)
  ld s1,OFFSET_REG_S1(tp)
  //args
  ld a0,OFFSET_REG_A0(tp)
  ld a1,OFFSET_REG_A1(tp)
  ld a2,OFFSET_REG_A2(tp)
  ld a3,OFFSET_REG_A3(tp)
  ld a4,OFFSET_REG_A4(tp)
  ld a5,OFFSET_REG_A5(tp)
  ld a6,OFFSET_REG_A6(tp)
  ld a7,OFFSET_REG_A7(tp)
  //saved registers
  ld s2,OFFSET_REG_S2(tp)
  ld s3,OFFSET_REG_S3(tp)
  ld s4,OFFSET_REG_S4(tp)
  ld s5,OFFSET_REG_S5(tp)
  ld s6,OFFSET_REG_S6(tp)
  ld s7,OFFSET_REG_S7(tp)
  ld s8,OFFSET_REG_S8(tp)
  ld s9,OFFSET_REG_S9(tp)
  ld s10,OFFSET_REG_S10(tp)
  ld s11,OFFSET_REG_S11(tp)
  //temporary 
  ld t3,OFFSET_REG_T3(tp)
  ld t4,OFFSET_REG_T4(tp)
  ld t5,OFFSET_REG_T5(tp)
  ld t6,OFFSET_REG_T6(tp)
  //privileged
  ld t3,OFFSET_REG_SSTATUS(tp)
 // csrw CSR_SSTATUS,t3
  csrw sstatus,t3
  csrr t3,sstatus
  ld t3,OFFSET_REG_SEPC(tp)
  //csrw CSR_SEPC,t3
  csrw sepc,t3
  csrr t3,sepc
  ld t3,OFFSET_REG_SBADADDR(tp)
  //csrw CSR_STVAL,t3
  csrw stval,t3
  ld t3,OFFSET_REG_SCAUSE(tp)
  //csrw CSR_SCAUSE,t3
  csrw scause,t3
  li t3,0xa22
  csrw sie,t3
  ld t3,OFFSET_REG_T3(tp)
.endm

ENTRY(enable_preempt)
  ld t1, current_running
  ld t0, PCB_PREEMPT_COUNT(t1)
  beq t0, zero, do_enable
  addi t0, t0, -1
  sd t0, PCB_PREEMPT_COUNT(t1)
  beq t0, zero, do_enable
  jr ra
do_enable:
  not t0, x0
  csrs CSR_SIE, t0
  jr ra
ENDPROC(enable_preempt)

ENTRY(disable_preempt)
  csrw CSR_SIE, zero
  ld t1, current_running
  ld t0, PCB_PREEMPT_COUNT(t1)
  addi t0, t0, 1
  sd t0, PCB_PREEMPT_COUNT(t1)
  jr ra
ENDPROC(disable_preempt)

ENTRY(enable_interrupt)
  li t0, SR_SIE
  csrw CSR_SSTATUS, t0
  jr ra
ENDPROC(enable_interrupt)

ENTRY(disable_interrupt)
  li t0, SR_SIE
  csrw CSR_SSTATUS, t0
  jr ra
ENDPROC(disable_interrupt)

// the address of previous pcb in a0
// the address of next pcb in a1
ENTRY(do_scheduler)
  rdtime a0
  sd a0,0(sp)
  jal scheduler
  lui t1,%hi(current_running)
  addi t1,t1,%lo(current_running)
  ld  tp,0(t1)
  rdtime a0
  ld a1,0(sp)
  sub a0,a0,a1
  //la a1,sched_cost
  //sd a0,0(a1)
  jal measure_cost
  j ret_from_exception
ENDPROC(do_scheduler)

ENTRY(ret_from_exception)
  /* TODO: */
  RESTORE_CONTEXT
  sret
ENDPROC(ret_from_exception)

ENTRY(exception_handler_entry)

  SAVE_CONTEXT
  csrw CSR_SSCRATCH, x0

  /* Load the global pointer */
  .option push
  .option norelax
  la gp, __global_pointer$
  .option pop

  /* TODO: load ret_from_exception into $ra
   * so that we can return to ret_from_exception
   * when interrupt_help complete.
   */
  lui ra,%hi(ret_from_exception)
  addi ra,ra,%lo(ret_from_exception)
  add a0,zero,tp
  csrr a2, scause
  csrr a1, CSR_STVAL
  /* TODO: call interrupt_helper
   * note: don not forget to pass parameters for it.
   */
  
  j interrupt_helper
  

ENDPROC(exception_handler_entry)
